package com.open.basic.util;

import com.open.basic.exception.SystemException;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.TimeUnit;


/**
 * @author: system
 * @Date: 2024/9/12 下午2:50
 * @Version: V1.0
 * @description: -> 单机版限流工具类
 */

@Slf4j
public class LimitUtil {

    private static final Map<String, RateLimiter> rateLimiterMap = Maps.newConcurrentMap();

    /**
     * 获取令牌
     *
     * @param key       限流key
     * @param rateLimit 每秒放入的令牌数
     */
    public static void acquire(String key, int rateLimit) {
        long start = System.currentTimeMillis();
        RateLimiter rateLimiter = null;
        if (rateLimiterMap.containsKey(key)) {
            rateLimiter = rateLimiterMap.get(key);
        } else {
            // 创建一个每秒放入rateLimit个令牌的RateLimiter
            rateLimiter = RateLimiter.create(rateLimit);
            rateLimiterMap.putIfAbsent(key, rateLimiter);
        }
        if (!rateLimiter.tryAcquire(100, TimeUnit.MILLISECONDS)) {
            throw new SystemException("服务繁忙，请稍后再试");
        }
    }
}
